{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ed82c3b9",
   "metadata": {},
   "source": [
    "## Trainable Interventions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3b71a4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Zhengxuan Wu\"\n",
    "__version__ = \"11/28/2023\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bce59d35",
   "metadata": {},
   "source": [
    "### Overview\n",
    "\n",
    "When interventions are static, you are mannually look for interesting counterfactual behaviors. What if interventions are trainable? and what if you train interventions with certain counterfactual behaviors?\n",
    "\n",
    "We think, if you can train such interventions, you find a systematic way of affecting the causal circuits realized in the NNs. With certain types of interventions, e.g., basis respect ones (vanilla causal abstraction or DAS), you are doing causal abstraction.\n",
    "\n",
    "In this tutorial, we show how you can train interventions with customized dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "714f9e7b",
   "metadata": {},
   "source": [
    "### Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c5dd0623",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-01-11 01:23:38,042] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    # This library is our indicator that the required installs\n",
    "    # need to be done.\n",
    "    import pyvene\n",
    "\n",
    "except ModuleNotFoundError:\n",
    "    !pip install git+https://github.com/frankaging/pyvene.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a5859137",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ffff1b0633df473ab9d15a22bba0f29d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/498M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1fa2101566342028d7d19dd74a8299f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/119 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded model\n",
      "\n",
      "A program that the model can solve (Cs are input tokens):\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       "  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"293.4pt\" height=\"286.2pt\" viewBox=\"0 0 293.4 286.2\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n",
       " <metadata>\n",
       "  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n",
       "   <cc:Work>\n",
       "    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n",
       "    <dc:date>2024-01-20T10:50:39.986017</dc:date>\n",
       "    <dc:format>image/svg+xml</dc:format>\n",
       "    <dc:creator>\n",
       "     <cc:Agent>\n",
       "      <dc:title>Matplotlib v3.7.3, https://matplotlib.org/</dc:title>\n",
       "     </cc:Agent>\n",
       "    </dc:creator>\n",
       "   </cc:Work>\n",
       "  </rdf:RDF>\n",
       " </metadata>\n",
       " <defs>\n",
       "  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n",
       " </defs>\n",
       " <g id=\"figure_1\">\n",
       "  <g id=\"patch_1\">\n",
       "   <path d=\"M 0 286.2 \n",
       "L 293.4 286.2 \n",
       "L 293.4 0 \n",
       "L 0 0 \n",
       "L 0 286.2 \n",
       "z\n",
       "\" style=\"fill: none\"/>\n",
       "  </g>\n",
       "  <g id=\"axes_1\">\n",
       "   <g id=\"patch_2\">\n",
       "    <path d=\"M 138.633892 242.911365 \n",
       "Q 93.861086 225.464414 51.17175 208.829343 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 57.173462 215.461044 \n",
       "L 51.17175 208.829343 \n",
       "L 60.078137 208.006993 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_3\">\n",
       "    <path d=\"M 139.872351 240.733384 \n",
       "Q 120.282172 225.465683 102.455689 211.572528 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 106.30682 219.645247 \n",
       "L 102.455689 211.572528 \n",
       "L 111.224543 213.335254 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_4\">\n",
       "    <path d=\"M 146.7 237.395405 \n",
       "Q 146.7 225.463409 146.7 215.767482 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 142.7 223.767482 \n",
       "L 146.7 215.767482 \n",
       "L 150.7 223.767482 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_5\">\n",
       "    <path d=\"M 153.527649 240.733384 \n",
       "Q 173.117828 225.465683 190.944311 211.572528 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 182.175457 213.335254 \n",
       "L 190.944311 211.572528 \n",
       "L 187.09318 219.645247 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_6\">\n",
       "    <path d=\"M 154.766108 242.911365 \n",
       "Q 199.538914 225.464414 242.22825 208.829343 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 233.321863 208.006993 \n",
       "L 242.22825 208.829343 \n",
       "L 236.226538 215.461044 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_7\">\n",
       "    <path d=\"M 48.702915 200.879971 \n",
       "Q 80.648967 184.281765 110.610792 168.714502 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 101.667599 168.853429 \n",
       "L 110.610792 168.714502 \n",
       "L 105.356017 175.952412 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_8\">\n",
       "    <path d=\"M 102.127577 202.295093 \n",
       "Q 159.912017 184.28126 215.561714 166.932916 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 206.733765 165.495098 \n",
       "L 215.561714 166.932916 \n",
       "L 209.11469 173.132584 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_9\">\n",
       "    <path d=\"M 139.015266 200.879971 \n",
       "Q 107.069215 184.281765 77.10739 168.714502 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 82.362164 175.952412 \n",
       "L 77.10739 168.714502 \n",
       "L 86.050583 168.853429 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_10\">\n",
       "    <path d=\"M 204.215769 197.585979 \n",
       "Q 212.751286 184.281585 220.079366 172.859236 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 212.3928 177.432714 \n",
       "L 220.079366 172.859236 \n",
       "L 219.12621 181.752575 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_11\">\n",
       "    <path d=\"M 243.929901 202.990715 \n",
       "Q 159.911 184.28199 78.07471 166.059274 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 85.014063 171.702445 \n",
       "L 78.07471 166.059274 \n",
       "L 86.752859 163.893695 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_12\">\n",
       "    <path d=\"M 244.113332 202.295093 \n",
       "Q 186.328892 184.28126 130.679195 166.932916 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 137.126219 173.132584 \n",
       "L 130.679195 166.932916 \n",
       "L 139.507144 165.495098 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_13\">\n",
       "    <path d=\"M 75.12337 159.698153 \n",
       "Q 107.069421 143.099946 137.031247 127.532684 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 128.088054 127.67161 \n",
       "L 137.031247 127.532684 \n",
       "L 131.776472 134.770594 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_14\">\n",
       "    <path d=\"M 124.954405 156.40416 \n",
       "Q 133.489923 143.099766 140.818002 131.677417 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 133.131437 136.250896 \n",
       "L 140.818002 131.677417 \n",
       "L 139.864846 140.570757 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_15\">\n",
       "    <path d=\"M 219.955322 157.449792 \n",
       "Q 186.331156 122.509584 154.257495 89.180566 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 156.922562 97.718564 \n",
       "L 154.257495 89.180566 \n",
       "L 162.686934 92.171311 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_16\">\n",
       "    <path d=\"M 146.7 113.84995 \n",
       "Q 146.7 101.917955 146.7 92.222027 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 142.7 100.222027 \n",
       "L 146.7 92.222027 \n",
       "L 150.7 100.222027 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"patch_17\">\n",
       "    <path d=\"M 146.7 72.668132 \n",
       "Q 146.7 60.736136 146.7 51.040209 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "    <path d=\"M 142.7 59.040209 \n",
       "L 146.7 51.040209 \n",
       "L 150.7 59.040209 \n",
       "\" clip-path=\"url(#p1488db3fe2)\" style=\"fill: none; stroke: #0000ff; stroke-opacity: 0.5; stroke-width: 2; stroke-linecap: round\"/>\n",
       "   </g>\n",
       "   <g id=\"PathCollection_1\">\n",
       "    <defs>\n",
       "     <path id=\"C0_0_20a4b1a3cd\" d=\"M 0 14.142136 \n",
       "C 3.750538 14.142136 7.347969 12.652031 10 10 \n",
       "C 12.652031 7.347969 14.142136 3.750538 14.142136 -0 \n",
       "C 14.142136 -3.750538 12.652031 -7.347969 10 -10 \n",
       "C 7.347969 -12.652031 3.750538 -14.142136 0 -14.142136 \n",
       "C -3.750538 -14.142136 -7.347969 -12.652031 -10 -10 \n",
       "C -12.652031 -7.347969 -14.142136 -3.750538 -14.142136 0 \n",
       "C -14.142136 3.750538 -12.652031 7.347969 -10 10 \n",
       "C -7.347969 12.652031 -3.750538 14.142136 0 14.142136 \n",
       "z\n",
       "\"/>\n",
       "    </defs>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"146.7\" y=\"246.054545\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"41.018182\" y=\"204.872727\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"93.859091\" y=\"204.872727\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"146.7\" y=\"204.872727\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"199.540909\" y=\"204.872727\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"252.381818\" y=\"204.872727\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"67.438636\" y=\"163.690909\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"120.279545\" y=\"163.690909\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"225.961364\" y=\"163.690909\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"146.7\" y=\"122.509091\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"146.7\" y=\"81.327273\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <use xlink:href=\"#C0_0_20a4b1a3cd\" x=\"146.7\" y=\"40.145455\" style=\"fill: #0000ff; fill-opacity: 0.2; stroke: #0000ff; stroke-opacity: 0.2\"/>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_1\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- IN -->\n",
       "     <g transform=\"translate(140.44125 249.365795) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-49\" d=\"M 628 4666 \n",
       "L 1259 4666 \n",
       "L 1259 0 \n",
       "L 628 0 \n",
       "L 628 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "       <path id=\"DejaVuSans-4e\" d=\"M 628 4666 \n",
       "L 1478 4666 \n",
       "L 3547 763 \n",
       "L 3547 4666 \n",
       "L 4159 4666 \n",
       "L 4159 0 \n",
       "L 3309 0 \n",
       "L 1241 3903 \n",
       "L 1241 0 \n",
       "L 628 0 \n",
       "L 628 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-49\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-4e\" x=\"29.492188\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_2\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- C0 -->\n",
       "     <g transform=\"translate(33.010994 208.183977) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-43\" d=\"M 4122 4306 \n",
       "L 4122 3641 \n",
       "Q 3803 3938 3442 4084 \n",
       "Q 3081 4231 2675 4231 \n",
       "Q 1875 4231 1450 3742 \n",
       "Q 1025 3253 1025 2328 \n",
       "Q 1025 1406 1450 917 \n",
       "Q 1875 428 2675 428 \n",
       "Q 3081 428 3442 575 \n",
       "Q 3803 722 4122 1019 \n",
       "L 4122 359 \n",
       "Q 3791 134 3420 21 \n",
       "Q 3050 -91 2638 -91 \n",
       "Q 1578 -91 968 557 \n",
       "Q 359 1206 359 2328 \n",
       "Q 359 3453 968 4101 \n",
       "Q 1578 4750 2638 4750 \n",
       "Q 3056 4750 3426 4639 \n",
       "Q 3797 4528 4122 4306 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "       <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \n",
       "Q 1547 4250 1301 3770 \n",
       "Q 1056 3291 1056 2328 \n",
       "Q 1056 1369 1301 889 \n",
       "Q 1547 409 2034 409 \n",
       "Q 2525 409 2770 889 \n",
       "Q 3016 1369 3016 2328 \n",
       "Q 3016 3291 2770 3770 \n",
       "Q 2525 4250 2034 4250 \n",
       "z\n",
       "M 2034 4750 \n",
       "Q 2819 4750 3233 4129 \n",
       "Q 3647 3509 3647 2328 \n",
       "Q 3647 1150 3233 529 \n",
       "Q 2819 -91 2034 -91 \n",
       "Q 1250 -91 836 529 \n",
       "Q 422 1150 422 2328 \n",
       "Q 422 3509 836 4129 \n",
       "Q 1250 4750 2034 4750 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-43\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-30\" x=\"69.824219\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_3\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- C1 -->\n",
       "     <g transform=\"translate(85.851903 208.183977) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-31\" d=\"M 794 531 \n",
       "L 1825 531 \n",
       "L 1825 4091 \n",
       "L 703 3866 \n",
       "L 703 4441 \n",
       "L 1819 4666 \n",
       "L 2450 4666 \n",
       "L 2450 531 \n",
       "L 3481 531 \n",
       "L 3481 0 \n",
       "L 794 0 \n",
       "L 794 531 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-43\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-31\" x=\"69.824219\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_4\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- C2 -->\n",
       "     <g transform=\"translate(138.692813 208.183977) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-32\" d=\"M 1228 531 \n",
       "L 3431 531 \n",
       "L 3431 0 \n",
       "L 469 0 \n",
       "L 469 531 \n",
       "Q 828 903 1448 1529 \n",
       "Q 2069 2156 2228 2338 \n",
       "Q 2531 2678 2651 2914 \n",
       "Q 2772 3150 2772 3378 \n",
       "Q 2772 3750 2511 3984 \n",
       "Q 2250 4219 1831 4219 \n",
       "Q 1534 4219 1204 4116 \n",
       "Q 875 4013 500 3803 \n",
       "L 500 4441 \n",
       "Q 881 4594 1212 4672 \n",
       "Q 1544 4750 1819 4750 \n",
       "Q 2544 4750 2975 4387 \n",
       "Q 3406 4025 3406 3419 \n",
       "Q 3406 3131 3298 2873 \n",
       "Q 3191 2616 2906 2266 \n",
       "Q 2828 2175 2409 1742 \n",
       "Q 1991 1309 1228 531 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-43\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-32\" x=\"69.824219\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_5\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- C3 -->\n",
       "     <g transform=\"translate(191.533722 208.183977) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-33\" d=\"M 2597 2516 \n",
       "Q 3050 2419 3304 2112 \n",
       "Q 3559 1806 3559 1356 \n",
       "Q 3559 666 3084 287 \n",
       "Q 2609 -91 1734 -91 \n",
       "Q 1441 -91 1130 -33 \n",
       "Q 819 25 488 141 \n",
       "L 488 750 \n",
       "Q 750 597 1062 519 \n",
       "Q 1375 441 1716 441 \n",
       "Q 2309 441 2620 675 \n",
       "Q 2931 909 2931 1356 \n",
       "Q 2931 1769 2642 2001 \n",
       "Q 2353 2234 1838 2234 \n",
       "L 1294 2234 \n",
       "L 1294 2753 \n",
       "L 1863 2753 \n",
       "Q 2328 2753 2575 2939 \n",
       "Q 2822 3125 2822 3475 \n",
       "Q 2822 3834 2567 4026 \n",
       "Q 2313 4219 1838 4219 \n",
       "Q 1578 4219 1281 4162 \n",
       "Q 984 4106 628 3988 \n",
       "L 628 4550 \n",
       "Q 988 4650 1302 4700 \n",
       "Q 1616 4750 1894 4750 \n",
       "Q 2613 4750 3031 4423 \n",
       "Q 3450 4097 3450 3541 \n",
       "Q 3450 3153 3228 2886 \n",
       "Q 3006 2619 2597 2516 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-43\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-33\" x=\"69.824219\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_6\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- C4 -->\n",
       "     <g transform=\"translate(244.374631 208.183977) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \n",
       "L 825 1625 \n",
       "L 2419 1625 \n",
       "L 2419 4116 \n",
       "z\n",
       "M 2253 4666 \n",
       "L 3047 4666 \n",
       "L 3047 1625 \n",
       "L 3713 1625 \n",
       "L 3713 1100 \n",
       "L 3047 1100 \n",
       "L 3047 0 \n",
       "L 2419 0 \n",
       "L 2419 1100 \n",
       "L 313 1100 \n",
       "L 313 1709 \n",
       "L 2253 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-43\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-34\" x=\"69.824219\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_7\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- == -->\n",
       "     <g transform=\"translate(57.383011 167.002159) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-3d\" d=\"M 678 2906 \n",
       "L 4684 2906 \n",
       "L 4684 2381 \n",
       "L 678 2381 \n",
       "L 678 2906 \n",
       "z\n",
       "M 678 1631 \n",
       "L 4684 1631 \n",
       "L 4684 1100 \n",
       "L 678 1100 \n",
       "L 678 1631 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-3d\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-3d\" x=\"83.789062\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_8\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- != -->\n",
       "     <g transform=\"translate(112.846108 167.002159) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-21\" d=\"M 966 794 \n",
       "L 1600 794 \n",
       "L 1600 0 \n",
       "L 966 0 \n",
       "L 966 794 \n",
       "z\n",
       "M 966 4666 \n",
       "L 1600 4666 \n",
       "L 1600 2619 \n",
       "L 1538 1503 \n",
       "L 1031 1503 \n",
       "L 966 2619 \n",
       "L 966 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-21\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-3d\" x=\"40.087891\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_9\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- s -->\n",
       "     <g transform=\"translate(222.835739 167.002159) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \n",
       "L 2834 2853 \n",
       "Q 2591 2978 2328 3040 \n",
       "Q 2066 3103 1784 3103 \n",
       "Q 1356 3103 1142 2972 \n",
       "Q 928 2841 928 2578 \n",
       "Q 928 2378 1081 2264 \n",
       "Q 1234 2150 1697 2047 \n",
       "L 1894 2003 \n",
       "Q 2506 1872 2764 1633 \n",
       "Q 3022 1394 3022 966 \n",
       "Q 3022 478 2636 193 \n",
       "Q 2250 -91 1575 -91 \n",
       "Q 1294 -91 989 -36 \n",
       "Q 684 19 347 128 \n",
       "L 347 722 \n",
       "Q 666 556 975 473 \n",
       "Q 1284 391 1588 391 \n",
       "Q 1994 391 2212 530 \n",
       "Q 2431 669 2431 922 \n",
       "Q 2431 1156 2273 1281 \n",
       "Q 2116 1406 1581 1522 \n",
       "L 1381 1569 \n",
       "Q 847 1681 609 1914 \n",
       "Q 372 2147 372 2553 \n",
       "Q 372 3047 722 3315 \n",
       "Q 1072 3584 1716 3584 \n",
       "Q 2034 3584 2315 3537 \n",
       "Q 2597 3491 2834 3397 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-73\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_10\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- AND -->\n",
       "     <g transform=\"translate(133.486875 125.820341) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-41\" d=\"M 2188 4044 \n",
       "L 1331 1722 \n",
       "L 3047 1722 \n",
       "L 2188 4044 \n",
       "z\n",
       "M 1831 4666 \n",
       "L 2547 4666 \n",
       "L 4325 0 \n",
       "L 3669 0 \n",
       "L 3244 1197 \n",
       "L 1141 1197 \n",
       "L 716 0 \n",
       "L 50 0 \n",
       "L 1831 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "       <path id=\"DejaVuSans-44\" d=\"M 1259 4147 \n",
       "L 1259 519 \n",
       "L 2022 519 \n",
       "Q 2988 519 3436 956 \n",
       "Q 3884 1394 3884 2338 \n",
       "Q 3884 3275 3436 3711 \n",
       "Q 2988 4147 2022 4147 \n",
       "L 1259 4147 \n",
       "z\n",
       "M 628 4666 \n",
       "L 1925 4666 \n",
       "Q 3281 4666 3915 4102 \n",
       "Q 4550 3538 4550 2338 \n",
       "Q 4550 1131 3912 565 \n",
       "Q 3275 0 1925 0 \n",
       "L 628 0 \n",
       "L 628 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-41\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-4e\" x=\"68.408203\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-44\" x=\"143.212891\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_11\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- OR -->\n",
       "     <g transform=\"translate(137.807812 84.638523) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-4f\" d=\"M 2522 4238 \n",
       "Q 1834 4238 1429 3725 \n",
       "Q 1025 3213 1025 2328 \n",
       "Q 1025 1447 1429 934 \n",
       "Q 1834 422 2522 422 \n",
       "Q 3209 422 3611 934 \n",
       "Q 4013 1447 4013 2328 \n",
       "Q 4013 3213 3611 3725 \n",
       "Q 3209 4238 2522 4238 \n",
       "z\n",
       "M 2522 4750 \n",
       "Q 3503 4750 4090 4092 \n",
       "Q 4678 3434 4678 2328 \n",
       "Q 4678 1225 4090 567 \n",
       "Q 3503 -91 2522 -91 \n",
       "Q 1538 -91 948 565 \n",
       "Q 359 1222 359 2328 \n",
       "Q 359 3434 948 4092 \n",
       "Q 1538 4750 2522 4750 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "       <path id=\"DejaVuSans-52\" d=\"M 2841 2188 \n",
       "Q 3044 2119 3236 1894 \n",
       "Q 3428 1669 3622 1275 \n",
       "L 4263 0 \n",
       "L 3584 0 \n",
       "L 2988 1197 \n",
       "Q 2756 1666 2539 1819 \n",
       "Q 2322 1972 1947 1972 \n",
       "L 1259 1972 \n",
       "L 1259 0 \n",
       "L 628 0 \n",
       "L 628 4666 \n",
       "L 2053 4666 \n",
       "Q 2853 4666 3247 4331 \n",
       "Q 3641 3997 3641 3322 \n",
       "Q 3641 2881 3436 2590 \n",
       "Q 3231 2300 2841 2188 \n",
       "z\n",
       "M 1259 4147 \n",
       "L 1259 2491 \n",
       "L 2053 2491 \n",
       "Q 2509 2491 2742 2702 \n",
       "Q 2975 2913 2975 3322 \n",
       "Q 2975 3731 2742 3939 \n",
       "Q 2509 4147 2053 4147 \n",
       "L 1259 4147 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-4f\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-52\" x=\"78.710938\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "   <g id=\"text_12\">\n",
       "    <g clip-path=\"url(#p1488db3fe2)\">\n",
       "     <!-- OUT -->\n",
       "     <g transform=\"translate(133.920937 43.456705) scale(0.12 -0.12)\">\n",
       "      <defs>\n",
       "       <path id=\"DejaVuSans-55\" d=\"M 556 4666 \n",
       "L 1191 4666 \n",
       "L 1191 1831 \n",
       "Q 1191 1081 1462 751 \n",
       "Q 1734 422 2344 422 \n",
       "Q 2950 422 3222 751 \n",
       "Q 3494 1081 3494 1831 \n",
       "L 3494 4666 \n",
       "L 4128 4666 \n",
       "L 4128 1753 \n",
       "Q 4128 841 3676 375 \n",
       "Q 3225 -91 2344 -91 \n",
       "Q 1459 -91 1007 375 \n",
       "Q 556 841 556 1753 \n",
       "L 556 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "       <path id=\"DejaVuSans-54\" d=\"M -19 4666 \n",
       "L 3928 4666 \n",
       "L 3928 4134 \n",
       "L 2272 4134 \n",
       "L 2272 0 \n",
       "L 1638 0 \n",
       "L 1638 4134 \n",
       "L -19 4134 \n",
       "L -19 4666 \n",
       "z\n",
       "\" transform=\"scale(0.015625)\"/>\n",
       "      </defs>\n",
       "      <use xlink:href=\"#DejaVuSans-4f\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-55\" x=\"78.710938\"/>\n",
       "      <use xlink:href=\"#DejaVuSans-54\" x=\"151.904297\"/>\n",
       "     </g>\n",
       "    </g>\n",
       "   </g>\n",
       "  </g>\n",
       " </g>\n",
       " <defs>\n",
       "  <clipPath id=\"p1488db3fe2\">\n",
       "   <rect x=\"7.2\" y=\"7.2\" width=\"279\" height=\"271.8\"/>\n",
       "  </clipPath>\n",
       " </defs>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os, json, torch\n",
    "import numpy as np\n",
    "np.object = object\n",
    "import pandas as pd\n",
    "from torch.utils.data import DataLoader\n",
    "from datasets.utils.logging import disable_progress_bar\n",
    "disable_progress_bar()\n",
    "\n",
    "from pyvene import embed_to_distrib, top_vals, format_token\n",
    "from pyvene import (\n",
    "    IntervenableModel,\n",
    "    LowRankRotatedSpaceIntervention,\n",
    "    RepresentationConfig,\n",
    "    IntervenableConfig,\n",
    ")\n",
    "from pyvene import create_gpt2_lm\n",
    "\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "from plotnine import (\n",
    "    ggplot,\n",
    "    geom_tile,\n",
    "    aes,\n",
    "    facet_wrap,\n",
    "    theme,\n",
    "    element_text,\n",
    "    geom_bar,\n",
    "    geom_hline,\n",
    "    scale_y_log10,\n",
    ")\n",
    "\n",
    "from tutorial_intervention_training_utils import (\n",
    "    visualize_program,\n",
    "    make_supervised_counterfactual_data_module,\n",
    ")\n",
    "from torch.nn import CrossEntropyLoss\n",
    "\n",
    "programs = json.load(open(os.path.join(\"./tutorial_data/\", \"seed_programs.json\")))\n",
    "programs_10 = json.load(\n",
    "    open(os.path.join(\"./tutorial_data/\", \"selected_programs_10.json\"))\n",
    ")\n",
    "\n",
    "# load fine-tuned model for this tutorial from HF\n",
    "config, tokenizer, gpt2 = create_gpt2_lm(\"zhengxuanzenwu/gpt2-5token-solver\")\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "_ = gpt2.to(\"cuda\")\n",
    "_ = gpt2.eval()\n",
    "\n",
    "# what is this model finetuned for? this visualizes a program that this model can solve.\n",
    "select_program = \"07065a\"\n",
    "print(\"\\nA program that the model can solve (Cs are input tokens):\")\n",
    "visualize_program(programs[select_program])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93011597",
   "metadata": {},
   "source": [
    "### Aligning the output variable with a single dimension in the rotated basis\n",
    "\n",
    "We know the output is a boolean value, which can be represented using a single dimension (i.e., pos/neg lies on a line with a linearly separable boundary). Can we learn a 1-d DAS direction to align activations with the output variable?\n",
    "\n",
    "We first need to sample (base, source(s)) pairs. In this case, it is easy, since the counterfactual behavior is that source output label overwrites the base one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "05f0f35b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model input looks like this (10-shot ICL):\n",
      "reportprint,itches,reenshots,Jean,reenshots=True\n",
      "dated,net,dated,Kate,dated=False\n",
      "Inst,Amid,Billy,While,Billy=True\n",
      "mas,quickShip,mas,imei,mas=True\n",
      "Excellent,embedreportprint,Ret,Theme,Excellent=True\n",
      "engers,Rick,engers,Jerry,engers=True\n",
      "debian,OSED,international,quickShip,international=True\n",
      "Led,orous,Lead,Cash,Lead=True\n",
      "mission,embedreportprint,mission,mbudsman,mission=True\n",
      "ABC,store,BBC,ixty,CBC=False\n",
      "ilitation,Brow,ilitation,ICLE,ilitation\n",
      "\n",
      "Training data for the intervention should contain these fields:\n",
      "dict_keys(['input_ids', 'attention_mask', 'source_input_ids', 'source_attention_mask', 'labels', 'source_0->base.0.pos', 'source_0->base.1.pos', 'subspaces'])\n"
     ]
    }
   ],
   "source": [
    "# load in some datasets we need\n",
    "num_of_shots = 10\n",
    "mode = \"E\"\n",
    "counterfactual_data_module, _ = make_supervised_counterfactual_data_module(\n",
    "    programs[select_program],\n",
    "    800,\n",
    "    num_of_shots,\n",
    "    \"op5\",  # this is the last variable = output\n",
    "    tokenizer,\n",
    "    data_path=\"./tutorial_data/\",\n",
    "    program_uuid=select_program,\n",
    "    mode=mode,\n",
    "    n_test_sample=1000,\n",
    ")\n",
    "\n",
    "print(\"Model input looks like this (10-shot ICL):\")\n",
    "print(counterfactual_data_module[\"train_dataset\"][0][\"question\"])\n",
    "\n",
    "# dataset cleaning\n",
    "columns_to_remove = [\n",
    "    \"question\",\n",
    "    \"source_question\",\n",
    "    \"answers\",\n",
    "    \"base_answers\",\n",
    "    \"source_answers\",\n",
    "]\n",
    "train_dataset = (\n",
    "    counterfactual_data_module[\"train_dataset\"]\n",
    "    .remove_columns(columns_to_remove)\n",
    "    .with_format(\"torch\")\n",
    ")\n",
    "eval_dataset = (\n",
    "    counterfactual_data_module[\"eval_dataset\"]\n",
    "    .remove_columns(columns_to_remove)\n",
    "    .with_format(\"torch\")\n",
    ")\n",
    "test_dataset = (\n",
    "    counterfactual_data_module[\"test_dataset\"]\n",
    "    .remove_columns(columns_to_remove)\n",
    "    .with_format(\"torch\")\n",
    ")\n",
    "\n",
    "\n",
    "def add_locations(example):\n",
    "    example[\"source_0->base.0.pos\"] = [129]  # the fixed last token location\n",
    "    example[\"source_0->base.1.pos\"] = [129]  # the fixed last token location\n",
    "    example[\"subspaces\"] = [0]  # the only subspace is a single column\n",
    "    return example\n",
    "\n",
    "\n",
    "train_dataset = train_dataset.map(add_locations).shuffle(seed=42)\n",
    "eval_dataset = eval_dataset.map(add_locations)\n",
    "test_dataset = test_dataset.map(add_locations)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64)\n",
    "eval_dataloader = DataLoader(eval_dataset, batch_size=64)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64)\n",
    "\n",
    "print(\"\\nTraining data for the intervention should contain these fields:\")\n",
    "print(train_dataset[0].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd7cbc9a",
   "metadata": {},
   "source": [
    "### A 1-d DAS direction learning config\n",
    "\n",
    "No full rotation learning is needed here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "55c5687a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_d_low_rank_das_position_config(\n",
    "    model_type, intervention_type, layer, intervention_types\n",
    "):\n",
    "    config = IntervenableConfig(\n",
    "        model_type=model_type,\n",
    "        representations=[\n",
    "            RepresentationConfig(\n",
    "                layer,              # layer\n",
    "                intervention_type,  # intervention type\n",
    "                \"pos\",              # intervention unit\n",
    "                1,                  # max number of unit\n",
    "                low_rank_dimension=1,         # a single das direction\n",
    "                subspace_partition=[[0, 1]],  # dummy partition\n",
    "            ),\n",
    "        ],\n",
    "        intervention_types=intervention_types,\n",
    "    )\n",
    "    return config\n",
    "\n",
    "\n",
    "config = single_d_low_rank_das_position_config(\n",
    "    type(gpt2), \"block_output\", 11, LowRankRotatedSpaceIntervention\n",
    ")\n",
    "intervenable = IntervenableModel(config, gpt2)\n",
    "intervenable.set_device(\"cuda\")\n",
    "intervenable.disable_model_gradients()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7654d6a",
   "metadata": {},
   "source": [
    "### Your own loss and metrics function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5cb7919d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inputs_collator(inputs):\n",
    "    for k, v in inputs.items():\n",
    "        if \"->\" in k:\n",
    "            inputs[k] = v.tolist()\n",
    "        elif \"subspace\" in k:\n",
    "            inputs[k] = [v.tolist()]\n",
    "        elif v is not None and isinstance(v, torch.Tensor):\n",
    "            inputs[k] = v.to(\"cuda\")\n",
    "    return inputs\n",
    "\n",
    "\n",
    "def compute_loss(logits, labels):\n",
    "    shift_logits = logits[..., :-1, :].contiguous()\n",
    "    shift_labels = labels[..., 1:].contiguous()\n",
    "\n",
    "    # Flatten the tokens\n",
    "    loss_fct = CrossEntropyLoss()\n",
    "    shift_logits = shift_logits.view(-1, 50257)\n",
    "    shift_labels = shift_labels.view(-1)\n",
    "    # Enable model parallelism\n",
    "    shift_labels = shift_labels.to(shift_logits.device)\n",
    "    loss = loss_fct(shift_logits, shift_labels)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def compute_metrics(eval_preds, eval_labels):\n",
    "    total_count = 0\n",
    "    correct_count = 0\n",
    "    for eval_pred, eval_label in zip(eval_preds, eval_labels):\n",
    "        actual_test_labels = eval_label[:, -1]\n",
    "        pred_test_labels = torch.argmax(eval_pred[:, -2], dim=-1)\n",
    "        correct_labels = actual_test_labels == pred_test_labels\n",
    "        total_count += len(correct_labels)\n",
    "        correct_count += correct_labels.sum().tolist()\n",
    "    accuracy = round(correct_count / total_count, 2)\n",
    "    return {\"accuracy\": accuracy}[\"accuracy\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa940fb2",
   "metadata": {},
   "source": [
    "### Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b78405cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [01:34<00:00,  9.43s/it, loss=0.24, acc=0.91]\n"
     ]
    }
   ],
   "source": [
    "intervenable.train_alignment(\n",
    "    train_dataloader=train_dataloader,\n",
    "    compute_loss=compute_loss,\n",
    "    compute_metrics=compute_metrics,\n",
    "    inputs_collator=inputs_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32581404",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "05a7b187",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                                                                                                                                                                          \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.9763200000000002"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intervenable.eval_alignment(\n",
    "    eval_dataloader=test_dataloader,\n",
    "    compute_metrics=compute_metrics,\n",
    "    inputs_collator=inputs_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec8e0fa8",
   "metadata": {},
   "source": [
    "The above means >97% of time, you can intervene on a single activation in the rotated basis to flip the model prediction"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
